"""
无人机通信仿真器 - 主要仿真逻辑
"""

import time
import random
import numpy as np
import config
from uav import UAV, UAVState
from message import Message, EmergencyMessage
from protocol import (
    TraditionalProtocol, 
    ZigBeeProtocol, 
    WiFiDirectProtocol, 
    LoRaProtocol, 
    FANETProtocol, 
    DCDPProtocol
)
import math

class UAVCommunicationSimulator:
    """无人机通信仿真器类"""
    
    def __init__(self):
        """初始化仿真器"""
        # 设置随机种子
        np.random.seed(config.RANDOM_SEED)
        random.seed(config.RANDOM_SEED)
        
        # 创建无人机
        self.uavs = [UAV(i) for i in range(config.NUM_UAVS)]
        
        # 创建协议实例
        self.protocols = {
            'traditional': TraditionalProtocol(),
            'zigbee': ZigBeeProtocol(),
            'wifi_direct': WiFiDirectProtocol(),
            'lora': LoRaProtocol(),
            'fanet': FANETProtocol(),
            'dcdp_protocol': DCDPProtocol()
        }
        
        # 传统协议组 - 用于比较分析
        self.traditional_protocols = ['traditional', 'zigbee', 'wifi_direct', 'lora', 'fanet']
        
        # 保留向后兼容性的引用
        self.traditional_protocol = self.protocols['traditional']
        self.dcdp_protocol = self.protocols['dcdp_protocol']
        
        # 仿真参数
        self.current_time = 0
        self.end_time = config.SIMULATION_TIME
        self.step_time = config.STEP_TIME
        
        # 仿真统计
        self.statistics = {
            'traditional': self._create_empty_stats(),
            'zigbee': self._create_empty_stats(),
            'wifi_direct': self._create_empty_stats(),
            'lora': self._create_empty_stats(),
            'fanet': self._create_empty_stats(),
            'dcdp_protocol': self._create_empty_stats()
        }
        
        # 消息生成参数
        self.message_generation_rate = 1.0  # 每秒每架无人机平均生成的消息数
        
        # 仿真事件记录
        self.events = []
        
        # 协议特性参数
        self.protocol_params = {
            'traditional': {
                'max_range': 500,        # 最大通信范围(米)
                'bandwidth': 0.5,        # 带宽(Mbps)
                'reliability': 0.7,      # 可靠性(成功率)
                'latency': 100           # 基础延迟(毫秒)
            },
            'zigbee': {
                'max_range': 100,        # 较短通信范围
                'bandwidth': 0.25,       # 低带宽
                'reliability': 0.85,     # 高可靠性
                'latency': 50            # 低延迟
            },
            'wifi_direct': {
                'max_range': 200,        # 中等通信范围
                'bandwidth': 5.0,        # 高带宽
                'reliability': 0.80,     # 中等可靠性
                'latency': 30            # 较低延迟
            },
            'lora': {
                'max_range': 1000,       # 长通信范围
                'bandwidth': 0.05,       # 极低带宽
                'reliability': 0.90,     # 高可靠性
                'latency': 150           # 高延迟
            },
            'fanet': {
                'max_range': 800,        # 较长通信范围
                'bandwidth': 1.0,        # 中等带宽
                'reliability': 0.75,     # 中等可靠性
                'latency': 80            # 中等延迟
            }
        }
    
    def _create_empty_stats(self):
        """创建空的统计数据结构"""
        return {
            'total_messages': 0,
            'successful_messages': 0,
            'lost_messages': 0,
            'emergency_messages_handled': 0,
            'avg_delay': 0,
            'missed_deadlines': 0,
            'emergency_messages': 0,
            'delays': []
        }
    
    def initialize(self, num_uavs=None):
        """
        初始化仿真环境
        
        参数:
            num_uavs: 无人机数量，如果为None则使用配置中的NUM_UAVS
        """
        # 设置无人机数量
        if num_uavs is not None:
            uav_count = num_uavs
        else:
            uav_count = config.NUM_UAVS
            
        # 创建无人机
        self.uavs = []
        for i in range(uav_count):
            # 随机生成UAV位置 (x, y, z)
            position = (
                random.uniform(0, config.AREA_SIZE),
                random.uniform(0, config.AREA_SIZE),
                random.uniform(50, 300)  # 高度在50-300米
            )
            
            # 随机速度和方向
            speed = random.uniform(5, 20)  # 5-20 m/s
            direction = (
                random.uniform(-1, 1),
                random.uniform(-1, 1),
                random.uniform(-0.2, 0.2)  # 小的垂直变化
            )
            
            # 标准化方向向量
            magnitude = math.sqrt(sum(d*d for d in direction))
            direction = tuple(d/magnitude for d in direction)
            
            # 创建UAV对象
            uav = {
                'id': i,
                'position': position,
                'speed': speed,
                'direction': direction,
                'battery': random.uniform(30, 100),  # 剩余电量百分比
                'task_priority': random.choice(['low', 'medium', 'high']),
                'type': random.choice(['surveillance', 'delivery', 'search', 'mapping'])
            }
            self.uavs.append(uav)
            
        # 更新无人机邻居关系
        self.update_uav_neighbors()
        
        print(f"初始化完成: 创建了{uav_count}架无人机")
    
    def update_uav_positions(self):
        """更新所有无人机的位置"""
        for uav in self.uavs:
            uav.update_position(self.step_time)
    
    def update_uav_neighbors(self):
        """更新所有无人机的邻居关系"""
        for uav in self.uavs:
            uav.update_neighbors(self.uavs)
    
    def generate_random_message(self, sender_id):
        """
        生成随机消息
        
        参数:
            sender_id: 发送方无人机ID
            
        返回:
            生成的消息
        """
        # 随机选择消息类型
        msg_types = list(config.MESSAGE_TYPES.keys())
        weights = [1, 3, 5, 2, 2]  # 各类型消息的相对权重
        msg_type = random.choices(msg_types, weights=weights, k=1)[0]
        
        # 随机选择接收方 (有10%概率是广播消息)
        if random.random() < 0.1:
            receiver_id = None  # 广播消息
        else:
            # 随机选择一个不是自己的无人机
            possible_receivers = [uav['id'] for uav in self.uavs if uav['id'] != sender_id]
            receiver_id = random.choice(possible_receivers)
        
        # 创建消息
        if msg_type == 'EMERGENCY':
            return EmergencyMessage(sender_id, receiver_id)
        else:
            return Message(sender_id, receiver_id, msg_type)
    
    def generate_messages(self):
        """为无人机生成随机消息"""
        for uav in self.uavs:
            # 决定是否生成消息
            if random.random() < self.message_generation_rate * self.step_time:
                message = self.generate_random_message(uav['id'])
                uav.add_message_to_queue(message)
                
                # 记录事件
                self.events.append({
                    'time': self.current_time,
                    'type': 'message_generated',
                    'uav_id': uav['id'],
                    'message': message
                })
    
    def process_messages(self):
        """处理所有无人机的消息队列"""
        # 随机排序无人机，以避免固定顺序带来的系统性偏差
        processed_uavs = list(self.uavs)
        random.shuffle(processed_uavs)
        
        for uav in processed_uavs:
            # 如果消息队列为空，跳过
            if not uav.message_queue:
                continue
            
            # 获取队首消息(优先级最高的)
            message = uav.message_queue[0]
            
            # 确定接收方
            if message.receiver_id is None:
                # 广播消息，发送给所有邻居
                for neighbor_id in uav.neighbors:
                    receiver = next((u for u in self.uavs if u['id'] == neighbor_id), None)
                    if receiver:
                        self.send_message(uav, receiver, message)
            else:
                # 点对点消息
                receiver = next((u for u in self.uavs if u['id'] == message.receiver_id), None)
                if receiver and receiver['id'] in uav.neighbors:
                    # 直接通信
                    self.send_message(uav, receiver, message)
                elif receiver:
                    # 需要路由
                    self.route_message(uav, receiver, message)
            
            # 移除队首消息
            uav.message_queue.pop(0)
    
    def send_message(self, sender, receiver, message):
        """
        发送消息并记录结果
        
        参数:
            sender: 发送方无人机
            receiver: 接收方无人机
            message: 要发送的消息
        """
        # 基于消息属性选择最合适的协议
        # 如果消息有指定协议属性，使用指定的协议
        specified_protocol = getattr(message, 'protocol', None)
        
        if specified_protocol and specified_protocol in self.protocols:
            protocol_name = specified_protocol
            protocol = self.protocols[protocol_name]
        else:
            # 随机选择协议，根据场景特性进行权重分配
            distance = np.linalg.norm(sender['position'] - receiver['position'])
            rel_speed = np.linalg.norm(sender['speed'] - receiver['speed'])
            is_emergency = message.type == 'EMERGENCY'
            
            # 基于场景特性为不同协议分配权重
            weights = {
                'traditional': 0.1,  # 基础协议作为基准
                'zigbee': 0.15 if distance < 100 else 0.05,  # ZigBee适合短距离
                'wifi_direct': 0.2 if distance < 200 else 0.05,  # WiFi Direct适合中等距离
                'lora': 0.3 if distance > 300 else 0.1,  # LoRa适合长距离
                'fanet': 0.2 if rel_speed > 5 else 0.1,  # FANET适合高移动性
                'dcdp_protocol': 0.5 if is_emergency else 0.3  # DCDP协议优先处理紧急消息
            }
            
            # 标准化权重
            total_weight = sum(weights.values())
            normalized_weights = {k: v/total_weight for k, v in weights.items()}
            
            # 随机选择协议
            protocols = list(normalized_weights.keys())
            weights_list = [normalized_weights[p] for p in protocols]
            protocol_name = random.choices(protocols, weights=weights_list, k=1)[0]
            protocol = self.protocols[protocol_name]
        
        # 发送消息
        success, delay = protocol.send_message(sender, receiver, message, self.current_time)
        
        # 更新统计
        self.statistics[protocol_name]['total_messages'] += 1
        if success:
            self.statistics[protocol_name]['successful_messages'] += 1
            # 更新平均延迟
            current_avg = self.statistics[protocol_name]['avg_delay']
            current_total = current_avg * (self.statistics[protocol_name]['successful_messages'] - 1)
            new_avg = (current_total + delay) / self.statistics[protocol_name]['successful_messages']
            self.statistics[protocol_name]['avg_delay'] = new_avg
            
            # 处理接收方接收消息
            receiver.process_received_message(message)
            
            # 如果是紧急消息
            if message.type == 'EMERGENCY':
                self.statistics[protocol_name]['emergency_messages'] += 1
                self.statistics[protocol_name]['emergency_messages_handled'] += 1
                
            # 检查是否超过截止时间
            if delay > message.deadline:
                self.statistics[protocol_name]['missed_deadlines'] += 1
        else:
            self.statistics[protocol_name]['lost_messages'] += 1
        
        # 记录事件
        self.events.append({
            'time': self.current_time,
            'type': 'message_sent',
            'protocol': protocol_name,
            'sender_id': sender['id'],
            'receiver_id': receiver['id'],
            'message_id': message.id,
            'message_type': message.type,
            'success': success,
            'delay': delay if success else None
        })
    
    def route_message(self, sender, target_receiver, message):
        """
        路由消息 (使用DCDP协议的路由功能)
        
        参数:
            sender: 发送方无人机
            target_receiver: 目标接收方无人机
            message: 要路由的消息
        """
        # 计算路由路径
        path = self.dcdp_protocol.route_message(self.uavs, message, self.current_time)
        
        # 如果找不到路径
        if not path or len(path) < 2:
            # 记录事件
            self.events.append({
                'time': self.current_time,
                'type': 'routing_failed',
                'sender_id': sender['id'],
                'target_id': target_receiver['id'],
                'message_id': message.id
            })
            return
        
        # 记录路由事件
        self.events.append({
            'time': self.current_time,
            'type': 'routing_started',
            'sender_id': sender['id'],
            'target_id': target_receiver['id'],
            'message_id': message.id,
            'path': path
        })
        
        # 实际上不进行路由传输，因为这里只关注直接通信性能比较
        # 真实路由实现需要考虑中继转发等复杂逻辑
    
    def inject_emergency(self, count=1):
        """
        注入紧急消息
        
        参数:
            count: 要注入的紧急消息数量
        """
        for _ in range(count):
            # 随机选择发送方
            sender = random.choice(self.uavs)
            
            # 创建紧急消息
            message = EmergencyMessage(
                sender_id=sender['id'],
                content={'emergency_type': 'CRITICAL', 'location': sender['position'].tolist()}
            )
            
            # 添加到发送队列
            sender.add_message_to_queue(message)
            
            # 记录事件
            self.events.append({
                'time': self.current_time,
                'type': 'emergency_injected',
                'uav_id': sender['id'],
                'message': message
            })
    
    def run_step(self):
        """运行单步仿真"""
        # 更新无人机位置
        self.update_uav_positions()
        
        # 更新邻居关系
        self.update_uav_neighbors()
        
        # 生成随机消息
        self.generate_messages()
        
        # 偶尔注入紧急消息
        if random.random() < 0.02:  # 2%的概率
            self.inject_emergency()
        
        # 处理消息队列
        self.process_messages()
        
        # 更新时间
        self.current_time += self.step_time
    
    def run_simulation(self, test_rounds=1000):
        """运行完整仿真"""
        print(f"开始无人机通信仿真 (总时长: {self.end_time}秒)...")
        
        real_start_time = time.time()
        
        # 主仿真循环
        while self.current_time < self.end_time:
            self.run_step()
            
            # 输出仿真进度
            if int(self.current_time) % 10 == 0 and int(self.current_time) % 10 == int((self.current_time - self.step_time) % 10) + self.step_time:
                progress = (self.current_time / self.end_time) * 100
                print(f"仿真进度: {progress:.1f}% (时间: {self.current_time:.1f}秒)")
        
        real_end_time = time.time()
        print(f"仿真完成! 实际耗时: {real_end_time - real_start_time:.2f}秒")
        
        # 返回仿真结果
        return {
            'statistics': self.statistics,
            'events': self.events,
            'uavs': [uav.get_status_summary() for uav in self.uavs],
            'protocols': {name: protocol.stats for name, protocol in self.protocols.items()}
        }
    
    def get_performance_comparison(self):
        """获取协议性能比较"""
        # 获取DCDP协议的统计数据
        dcdp = self.statistics['dcdp_protocol']
        
        # 聚合所有传统协议的统计数据
        traditional_combined = self._create_empty_stats()
        
        for protocol_name in self.traditional_protocols:
            protocol_stats = self.statistics[protocol_name]
            
            # 累加各项统计
            traditional_combined['total_messages'] += protocol_stats['total_messages']
            traditional_combined['successful_messages'] += protocol_stats['successful_messages']
            traditional_combined['lost_messages'] += protocol_stats['lost_messages']
            traditional_combined['emergency_messages'] += protocol_stats['emergency_messages']
            traditional_combined['emergency_messages_handled'] += protocol_stats['emergency_messages_handled']
            traditional_combined['missed_deadlines'] += protocol_stats['missed_deadlines']
        
        # 计算传统协议平均时延 (加权平均)
        weighted_delay_sum = 0
        total_success = 0
        
        for protocol_name in self.traditional_protocols:
            protocol_stats = self.statistics[protocol_name]
            successful = protocol_stats['successful_messages']
            if successful > 0:
                weighted_delay_sum += protocol_stats['avg_delay'] * successful
                total_success += successful
        
        if total_success > 0:
            traditional_combined['avg_delay'] = weighted_delay_sum / total_success
        
        # 计算成功率
        trad_success_rate = traditional_combined['successful_messages'] / max(1, traditional_combined['total_messages']) * 100
        dcdp_success_rate = dcdp['successful_messages'] / max(1, dcdp['total_messages']) * 100
        
        # 计算紧急消息处理率
        trad_emergency_rate = traditional_combined['emergency_messages_handled'] / max(1, traditional_combined['total_messages']) * 100
        dcdp_emergency_rate = dcdp['emergency_messages_handled'] / max(1, dcdp['total_messages']) * 100
        
        # 计算截止时间满足率
        trad_deadline_met = 100 - (traditional_combined['missed_deadlines'] / max(1, traditional_combined['successful_messages']) * 100)
        dcdp_deadline_met = 100 - (dcdp['missed_deadlines'] / max(1, dcdp['successful_messages']) * 100)
        
        return {
            'success_rate': {
                'traditional': trad_success_rate,
                'dcdp_protocol': dcdp_success_rate,
                'improvement': dcdp_success_rate - trad_success_rate
            },
            'avg_delay': {
                'traditional': traditional_combined['avg_delay'],
                'dcdp_protocol': dcdp['avg_delay'],
                'improvement': traditional_combined['avg_delay'] - dcdp['avg_delay']
            },
            'emergency_handling': {
                'traditional': trad_emergency_rate,
                'dcdp_protocol': dcdp_emergency_rate,
                'improvement': dcdp_emergency_rate - trad_emergency_rate
            },
            'deadline_met': {
                'traditional': trad_deadline_met,
                'dcdp_protocol': dcdp_deadline_met,
                'improvement': dcdp_deadline_met - trad_deadline_met
            }
        }
    
    def get_detailed_protocol_stats(self):
        """获取所有协议的详细统计信息"""
        results = {}
        
        for protocol_name, stats in self.statistics.items():
            if stats['total_messages'] == 0:
                continue
                
            success_rate = stats['successful_messages'] / max(1, stats['total_messages']) * 100
            emergency_rate = stats['emergency_messages'] / max(1, stats['total_messages']) * 100
            deadline_met = 100 - (stats['missed_deadlines'] / max(1, stats['successful_messages']) * 100)
            
            results[protocol_name] = {
                'success_rate': success_rate,
                'avg_delay': stats['avg_delay'],
                'emergency_handling': emergency_rate,
                'deadline_met': deadline_met,
                'total_messages': stats['total_messages']
            }
            
        return results 